
from utils.DorisQueryTool import DorisQueryTool
from TrainModel import trainModel
from Prediction import *
from utils.MysqlUtils import insert_data_to_mysql
import os

if __name__ == "__main__":
    # 查询所有收费站近两年数据

    with DorisQueryTool() as tool:

        df = tool.execute_query("SELECT count(1) flow,EX_TOLL_SECTION_ID roadId,EX_TOLL_STATION_ID stationId,EX_TOLL_STATION_NAME stationName,DATE_FORMAT(EX_TIME, '%Y-%m-%d') date FROM `toll_ex` where EX_TIME >= '2024-1-1 00:00:00'  group by EX_TOLL_SECTION_ID,EX_TOLL_STATION_ID,EX_TOLL_STATION_NAME ,date")
        result = df.sort_values('date').groupby('stationId', sort=False)

        for stationId, group in result:



            modelPath = "models/traffic_flow_model_" + stationId + ".pth"
            if os.path.exists(modelPath):
                continue
            # 每一个收费站训练一个模型，并保存
            trainModel(group,modelPath)
            # 使用刚预测好的模型预测未来7天数据
            # 加载模型
            model = load_pretrained_model(modelPath)

            # 示例：预测从"2023-12-01"开始的未来一周
            start_date = "2025-1-1"  # 可修改为实际日期
            end_date = "2025-6-9"  # 可修改为实际日期
            start_time = datetime.strptime(start_date, "%Y-%m-%d")
            end_time = datetime.strptime(end_date, "%Y-%m-%d")
            features_norm, scaler = preprocess_data(group)
            insert_datas = []
            while start_time <= end_time:
                prediction, dates = predict_future(model, scaler, start_time, group)
                # 7. 打印结果
                print("\n未来一周流量预测:")
                if (len(dates) > 0):
                    dict = {"roadId": group.iloc[0]["roadId"], "stationId": stationId, "flowPrediction": int(prediction[0]),
                            "datePrediction": dates[0], "createTime": datetime.now()}
                    # 查询真实车流量
                    target_row = group[group['date'] == dates[0]]
                    if not target_row.empty:
                        dict["flow"] = target_row['flow'].values[0]
                    else:
                        dict["flow"] = 0
                    insert_datas.append(dict)
                start_time = start_time + timedelta(days=1)

            # 执行入库操作
            insert_data_to_mysql(
                "INSERT INTO ex_station_prediction (road_id, station_id, flow_prediction,flow, date_prediction, create_time) VALUES (%(roadId)s, %(stationId)s, %(flowPrediction)s, %(flow)s, %(datePrediction)s, %(createTime)s);",
                insert_datas);

